from fastapi import FastAPI, WebSocket, Request, WebSocketDisconnect
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
import uvicorn
import json
import asyncio
import os
from datetime import datetime
import redis.asyncio as redis

app = FastAPI()

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
app.mount("/static", StaticFiles(directory=os.path.join(BASE_DIR, "static")), name="static")
templates = Jinja2Templates(directory=os.path.join(BASE_DIR, "templates"))

# Redis连接池
redis_pool = redis.ConnectionPool.from_url(
    "redis://:123456@localhost:6379",
    decode_responses=True,
    max_connections=100
)

# 全局连接管理器
class ConnectionManager:
    def __init__(self):
        self.redis = redis.Redis(connection_pool=redis_pool)

    async def add_connection(self, user_id: str, connection_id: str):
        """添加连接并同步到Redis"""
        await self.redis.hset("ws_connections", user_id, connection_id)
        await self.redis.sadd("online_users", user_id)

    async def remove_connection(self, user_id: str):
        """移除连接并同步到Redis"""
        await self.redis.hdel("ws_connections", user_id)
        await self.redis.srem("online_users", user_id)

    async def get_connection_id(self, user_id: str) -> str:
        """获取用户的连接ID"""
        return await self.redis.hget("ws_connections", user_id)

    async def is_user_online(self, user_id: str) -> bool:
        """检查用户是否在线"""
        return await self.redis.sismember("online_users", user_id)

    async def get_online_users(self):
        """获取所有在线用户"""
        return await self.redis.smembers("online_users")

manager = ConnectionManager()

@app.on_event("startup")
async def startup_event():
    # 清空 Redis 数据库（DB 0），生产环境请注释掉这行！
    await manager.redis.flushdb()
    asyncio.create_task(start_user_list_listener())

# 存储当前进程中的WebSocket连接映射 {connection_id: WebSocket}
current_process_connections = {}

@app.get("/")
async def get(request: Request):
    return templates.TemplateResponse("index.html", {"request": request})

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    user_id = None
    connection_id = str(id(websocket))
    
    try:
        # 1. 注册阶段
        data = await websocket.receive_text()
        message = json.loads(data)
        
        if message.get("type") != "register":
            await websocket.close(code=1003)
            return
            
        user_id = message["user_id"]
        
        # 原子注册
        registered = await manager.redis.setnx(f"user:{user_id}:lock", "1")
        if not registered:
            await websocket.close(code=1008)
            return
            
        try:
            # 存储用户信息
            await asyncio.gather(
                manager.add_connection(user_id, connection_id),
                manager.redis.hset("user_meta", user_id, json.dumps({
                    "id": user_id,
                    "name": message.get("name", user_id),
                    "avatar": message.get("avatar", f"https://api.dicebear.com/7.x/miniavs/svg?seed={user_id}"),
                    "status": "online"
                }))
            )
            
            # 将连接添加到当前进程的映射中
            current_process_connections[connection_id] = websocket

            # 2. 发送欢迎消息
            await websocket.send_text(json.dumps({
                "type": "system",
                "content": "Welcome to the chat!"
            }))

            # 3. 广播用户列表
            await broadcast_users()

            # 4. 创建消息监听器
            pubsub = manager.redis.pubsub()
            await pubsub.subscribe(f"user:{user_id}")

            async def message_listener():
                async for msg in pubsub.listen():
                    if msg["type"] == "message":
                        print("\n 1.msg[data]:\n", msg["data"])  # 保留原print
                        await websocket.send_text(msg["data"])

            listener_task = asyncio.create_task(message_listener())

            # 5. 主消息循环
            while True:
                data = await websocket.receive_text()
                msg = json.loads(data)
                
                if msg["type"] == "message":
                    target_id = msg["to"]
                    
                    if await manager.is_user_online(target_id):
                        await manager.redis.publish(
                            f"user:{target_id}",
                            json.dumps({
                                "type": "message",
                                "from": user_id,
                                "to": target_id,
                                "content": msg["content"],
                                "timestamp": msg.get("timestamp", datetime.now().isoformat())
                            })
                        )

        finally:
            await manager.redis.delete(f"user:{user_id}:lock")

    except WebSocketDisconnect:
        await handle_disconnect(user_id, connection_id)
    except Exception as e:
        print(f"Error: {e}")  # 保留原print
        await handle_disconnect(user_id, connection_id)

async def handle_disconnect(user_id: str, connection_id: str):
    # 从当前进程连接映射中移除
    if connection_id in current_process_connections:
        del current_process_connections[connection_id]
    
    # 从Redis中移除用户信息
    if user_id:
        await manager.remove_connection(user_id)
        await broadcast_users()

async def broadcast_users():
    """广播用户列表到所有进程"""
    # 获取所有在线用户
    online_users = await manager.get_online_users()
    print("\n online_users:\n", online_users)  # 保留原print

    # 新增：处理在线用户为空的情况
    if not online_users:
        print("警告：当前没有在线用户，跳过广播")
        return
    
    user_data = await manager.redis.hmget("user_meta", *online_users)

    print("\n user_data:\n", user_data)  # 保留原print

    users = []
    for data in user_data:
        if data:
            users.append(json.loads(data))
    
    print("\n users:\n", users)  # 保留原print

    message = json.dumps({
        "type": "users",
        "users": users
    })
    
    # 通过Redis发布用户列表更新
    await manager.redis.publish("user_list_update", message)

# 每个进程独立监听用户列表更新
async def start_user_list_listener():
    pubsub = manager.redis.pubsub()

    print("\n pubsub(start_user_list_listener):\n", pubsub)  # 保留原print

    await pubsub.subscribe("user_list_update")
    
    async for msg in pubsub.listen():
        if msg["type"] == "message":
            # 将消息发送给当前进程中的所有连接
            for connection_id, websocket in list(current_process_connections.items()):
                try:
                    print("\n 2.msg[data]:\n", msg["data"])  # 保留原print
                    await websocket.send_text(msg["data"])
                except:
                    # 连接可能已断开，从映射中移除
                    if connection_id in current_process_connections:
                        del current_process_connections[connection_id]

# 在应用启动时启动监听器（已合并重复定义）
@app.on_event("startup")
async def startup_event():
    asyncio.create_task(start_user_list_listener())

if __name__ == "__main__":
    uvicorn.run(
        "main:app",
        host="0.0.0.0",
        port=8000,
        workers=1,  # Windows下先使用单进程调试，稳定后再尝试多进程
        http="httptools",
        timeout_keep_alive=300
    )
